import torch
from torch.utils.data import Dataset, DataLoader

# 自定义 Dataset 类
class GPUDataset(Dataset):
    def __init__(self, data, labels):
        # 将数据和标签直接放在 GPU 上
        self.data = data.to('cuda')
        self.labels = labels.to('cuda')

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        # 返回已经在 GPU 上的数据
        return self.data[idx], self.labels[idx]

# 示例数据
data = torch.randn(1000, 3, 32, 32)  # 假设有 1000 个 3x32x32 的图像
labels = torch.randint(0, 10, (1000,))  # 假设有 1000 个标签

# 创建 GPU Dataset
dataset = GPUDataset(data, labels)

# 创建 DataLoader
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 训练循环示例
for batch_data, batch_labels in dataloader:
    # 由于数据已经在 GPU 上，这里不需要再调用 .to('cuda')
    print(batch_data.device)  # 应该输出: cuda:0
    print(batch_labels.device)  # 应该输出: cuda:0
    # 进行训练步骤...